task_flag="controlnetinpaint_1128L3data_style_load_dataset_norm"                                # the task flag is used to identify folders.
control_type=masked  # canny
resume=/llmcapagroup1/test-bucket/liuxin/hunyuantrainweightnew/L3train/018-datasetshuffle_batch2+batch3+batch4+oil7k+jiayouzhan194_size_multiscale_dit_g2_multiscale_L0base/checkpoints/0060000.pt/mp_rank_00_model_states.pt                               # checkpoint root for resume
# index_file=/llmcapagroup1/test-bucket/xinyu/code/HunyuanDIT-PRE-main/dataset/kolors_generated_51w_multiscale.json          # index file for dataloader
index_file=/llmcapagroup1/test-bucket/xinyu/code/HunyuanDIT-PRE-main/dataset/batch2+batch3+batch4+oil7k+jiayouzhan194_size_multiscale_bs4.json
results_dir=Log_train/inpainting_L3base                                       # save root for results
batch_size=4                                                 # training batch size
image_size=1024                                              # training image resolution
grad_accu_steps=1                                            # gradient accumulation
warmup_num_steps=0                                           # warm-up steps
lr=0.0001                                                    # learning rate  PYTHONPATH=./ sh hydit/train_controlnet.sh
ckpt_every=5000                                             # create a ckpt every a few steps.
ckpt_latest_every=5000                                       # create a ckpt named `latest.pt` every a few steps.


sh $(dirname "$0")/run_g_controlnet.sh \
    --task-flag ${task_flag} \
    --control-type ${control_type} \
    --noise-schedule scaled_linear --beta-start 0.00085 --beta-end 0.03 \
    --predict-type v_prediction \
    --reso-step 64 \
    --multireso \
    --uncond-p 0.05 \
    --uncond-p-t5 0.05 \
    --uncond-p-img 0.05 \
    --index-file ${index_file} \
    --random-flip \
    --lr ${lr} \
    --batch-size ${batch_size} \
    --image-size ${image_size} \
    --global-seed 999 \
    --grad-accu-steps ${grad_accu_steps} \
    --warmup-num-steps ${warmup_num_steps} \
    --use-flash-attn \
    --use-fp16 \
    --use-ema \
    --resume-ema-root ${resume} \
    --ema-dtype fp32 \
    --results-dir ${results_dir} \
    --reset-loader \
    --resume \
    --resume-module-root ${resume} \
    --ckpt-every ${ckpt_every} \
    --ckpt-latest-every ${ckpt_latest_every} \
    --log-every 10 \
    --deepspeed \
    --use-zero-stage 2 \
    --gradient-checkpointing \
    --no-strict \
    --use-style-cond \
    --size-cond 1024 1024 \
    --is-inpaint \
    "$@"
